package com.yf.service.chat.impl;

import cn.hutool.core.bean.BeanUtil;
import com.yf.exception.ServiceException;
import com.yf.model.ai.entity.AiConfig;
import com.yf.model.ai.entity.AiConversation;
import com.yf.model.ai.entity.AiMessage;
import com.yf.model.ai.entity.AiTool;
import com.yf.model.ai.enums.AiMessageRoleEnum;
import com.yf.model.ai.form.AiPromptForm;
import com.yf.result.ResultCode;
import com.yf.security.utils.SecurityUtil;
import com.yf.service.chat.IDeepSeekService;
import com.yf.service.manage.IAiConfigService;
import com.yf.service.manage.IAiConversationService;
import com.yf.service.manage.IAiConversationToolBindingService;
import com.yf.service.manage.IAiMessageService;
import com.yf.tool.context.ToolBaseContext;
import com.yf.tool.registrar.AiToolRegistrar;
import com.yf.utils.IpUtil;
import jakarta.servlet.http.HttpServletRequest;
import lombok.RequiredArgsConstructor;
import org.reactivestreams.Subscription;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.deepseek.DeepSeekChatModel;
import org.springframework.ai.deepseek.DeepSeekChatOptions;
import org.springframework.http.HttpHeaders;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

/**
 * AiService
 *
 * @author : YiFei
 * @since : 2024/6/7 9:30
 */
@Service("deepSeekService")
@RequiredArgsConstructor
public class DeepSeekServiceImpl implements IDeepSeekService {

    private final DeepSeekChatModel deepSeekChatModel;
    private final HttpServletRequest request;
    private final AiMessageChatMemory aiMessageChatMemory;
    private final IAiConversationService aiConversationService;
    private final IAiMessageService aiMessageService;
    private final IAiConfigService aiConfigService;
    private final AiToolRegistrar aiToolRegistrar;
    private final IAiConversationToolBindingService conversationToolBindingService;

    /**
     * 调用 DeepSeek 模型
     *
     * @param aiPromptForm 提问表单
     * @return 流式响应
     */
    @Override
    public Flux<ChatResponse> callStream(AiPromptForm aiPromptForm) {
        // 检查会话 ID 是否存在是否属于当前用户
        if (!checkConversationIdIsCurrentUser(aiPromptForm.getConversationId())) {
            return Flux.error(new ServiceException(ResultCode.AI_CONVERSATION_NOT_FOUND));
        }
        // 初始化模型配置
        DeepSeekChatOptions deepSeekChatOptions = initAiOptions(aiPromptForm.getConversationId());

        // 调用 DeepSeek 模型
        return ChatClient.builder(deepSeekChatModel)
                .defaultAdvisors(
                        MessageChatMemoryAdvisor.builder(aiMessageChatMemory)
                                .conversationId(String.valueOf(aiPromptForm.getConversationId()))
                                .build(),                                            // 会话历史记录
                        SimpleLoggerAdvisor.builder().build())                       // Debug 日志输出
                .build()
                .prompt(Prompt.builder()
                        .content(aiPromptForm.getPrompt())                           // 提问内容
                        .chatOptions(deepSeekChatOptions)                            // 模型选项
                        .build())
                .stream()
                .chatResponse();
    }

    /**
     * 重试AI输出
     *
     * @param conversationId 重试表单
     * @return 流式响应
     */
    @Override
    public Flux<ChatResponse> retry(Long conversationId) {
        if (!checkConversationIdIsCurrentUser(conversationId)) {
            return Flux.error(new ServiceException(ResultCode.AI_CONVERSATION_NOT_FOUND));
        }

        // 获取最后一条用户消息
        AiMessage aiMessage = aiMessageService.lambdaQuery()
                .eq(AiMessage::getConversationId, conversationId)
                .eq(AiMessage::getRole, AiMessageRoleEnum.USER)
                .orderByDesc(AiMessage::getCreateTime)
                .last("limit 1")
                .one();

        // 删除大于最后一条消息和之后的所有消息
        aiMessageService.lambdaUpdate()
                .eq(AiMessage::getConversationId, conversationId)
                .ge(AiMessage::getSendTime, aiMessage.getSendTime())
                .remove();

        // 重试请求
        return callStream(AiPromptForm.builder()
                .conversationId(conversationId)
                .prompt(aiMessage.getContent()) // 使用原消息内容进行重试
                .build());
    }

    /**
     * 删除消息
     *
     * @param conversationId 消息ID
     * @return 是否删除成功
     */
    @Override
    public Boolean deleteMessage(Long conversationId) {
        if (!checkConversationIdIsCurrentUser(conversationId)) {
            throw new ServiceException(ResultCode.AI_CONVERSATION_NOT_FOUND);
        }

        // 仅仅只能从最后一条消息开始删除
        aiMessageService.lambdaUpdate()
                .eq(AiMessage::getConversationId, conversationId)
                .orderByDesc(AiMessage::getCreateTime)
                .orderByDesc(AiMessage::getId) // 为了确保唯一性，加上主键降序
                .last("LIMIT 1")
                .remove();
        return true;
    }

    /**
     * 检查会话 ID 是否存在且属于当前用户
     *
     * @param conversationId 会话 ID
     * @return true 表示是，false 表示不是
     */
    private boolean checkConversationIdIsCurrentUser(Long conversationId) {
        if (conversationId == null) {
            return false; // 会话 ID 不能为空
        }
        // 检查会话是否存在且属于当前用户
        return aiConversationService.lambdaQuery()
                .eq(AiConversation::getId, conversationId)
                .eq(AiConversation::getUserId, SecurityUtil.getUserId())
                .exists();
    }

    /**
     * 初始化 DeepSeek 模型选项
     *
     * @param conversationId 会话 ID
     * @return DeepSeekChatOptions
     */
    private DeepSeekChatOptions initAiOptions(Long conversationId) {
        // 查询 ai 配置
        AiConfig aiConfig = aiConfigService.lambdaQuery()
                .eq(AiConfig::getConversationId, conversationId)
                .one();
        if (aiConfig == null) {
            // 如果未进行AI配置，则使用默认配置
            aiConfig = aiConfigService.lambdaQuery()
                    .eq(AiConfig::getIsDefault, true)
                    .one();
        }
        // 查询 ai 工具 , 区分不同类型工具
        List<AiTool> aiTools = conversationToolBindingService.getToolsByConversationId(conversationId);
        Set<String> aiToolNames = aiTools.stream().map(AiTool::getName).collect(Collectors.toSet());
        return DeepSeekChatOptions.builder()
                .model(aiConfig.getModel().getValue())                          // 模型名称
                .temperature(aiConfig.getTemperature())                         // 温度
                .toolContext(BeanUtil.beanToMap(initToolContext()))             // 工具类上下文
                .toolCallbacks(aiToolRegistrar.getToolCallbacks(aiToolNames))   // 获取所有需要的工具
                .build();
    }

    /**
     * 初始化工具所需基础上下文
     *
     * @return ToolContextDTO
     */
    private ToolBaseContext initToolContext() {
        return ToolBaseContext.builder()
                .userId(SecurityUtil.getUserId())
                .roles(SecurityUtil.getUserRoles())
                .ip(IpUtil.getIpAddr(request))
                .headerUserAgent(request.getHeader(HttpHeaders.USER_AGENT))
                .build();
    }

}
